import numpy as np
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt

spr = 1
spc = 1
spn = 0

T = [[3, 104, 98],
    [2, 100, 93],
    [1, 81, 95],
    [101, 10, 16],
    [99, 5, 8],
    [98, 2, 7]]

x = [18, 90]
k= 3

T = np.array(T, dtype=np.float64)
m = len(T)
x = np.array(x, dtype=np.float64)

spn += 1
ax = plt.subplot(spr, spc, spn, projection='3d')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.scatter3D(T[:, 0], T[:, 1], T[:, 2])


dis_vector = ((T[:, :2] - x)**2).sum(axis=1)**0.5
idx_sorted = dis_vector.argsort()
print(idx_sorted)

target_head = T[idx_sorted[:k], 2]
h = target_head.mean()
print(f'Hypothesis = {h}')
ax.scatter3D(x[0], x[1], h, color='r')

# by lib
from sklearn.neighbors import KNeighborsRegressor
model = KNeighborsRegressor(n_neighbors=k)
model.fit(T[:, :-1], T[:, -1])
# h = model.predict(x)
h = model.predict(x.reshape(1, -1))
print(f'Hypothesis by lib = {h}')

plt.show()
